"""
基于Vision Transformer的ES和扩展F信号分类模型
针对MX450显卡优化的轻量级实现
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from label_mapping import ID_TO_DESCRIPTION

class PatchEmbedding(nn.Module):
    """图像块嵌入层"""
    def __init__(self, img_size=224, patch_size=16, in_channels=3, embed_dim=768):
        super().__init__()
        self.img_size = img_size
        self.patch_size = patch_size
        self.n_patches = (img_size // patch_size) ** 2
        
        self.projection = nn.Conv2d(
            in_channels, embed_dim, 
            kernel_size=patch_size, stride=patch_size
        )
        
    def forward(self, x):
        # x: (batch_size, channels, height, width)
        x = self.projection(x)  # (batch_size, embed_dim, n_patches_sqrt, n_patches_sqrt)
        x = x.flatten(2)        # (batch_size, embed_dim, n_patches)
        x = x.transpose(1, 2)   # (batch_size, n_patches, embed_dim)
        return x

class MultiHeadAttention(nn.Module):
    """多头注意力机制"""
    def __init__(self, embed_dim=768, n_heads=12, dropout=0.1):
        super().__init__()
        self.embed_dim = embed_dim
        self.n_heads = n_heads
        self.head_dim = embed_dim // n_heads
        
        assert embed_dim % n_heads == 0, "embed_dim必须能被n_heads整除"
        
        self.qkv = nn.Linear(embed_dim, embed_dim * 3)
        self.proj = nn.Linear(embed_dim, embed_dim)
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, x):
        batch_size, seq_len, embed_dim = x.shape
        
        # 计算Q, K, V
        qkv = self.qkv(x).reshape(batch_size, seq_len, 3, self.n_heads, self.head_dim)
        qkv = qkv.permute(2, 0, 3, 1, 4)  # (3, batch_size, n_heads, seq_len, head_dim)
        q, k, v = qkv[0], qkv[1], qkv[2]
        
        # 计算注意力分数
        scale = math.sqrt(self.head_dim)
        attn = (q @ k.transpose(-2, -1)) / scale
        attn = F.softmax(attn, dim=-1)
        attn = self.dropout(attn)
        
        # 应用注意力
        out = attn @ v  # (batch_size, n_heads, seq_len, head_dim)
        out = out.transpose(1, 2).reshape(batch_size, seq_len, embed_dim)
        
        return self.proj(out)

class MLP(nn.Module):
    """前馈神经网络"""
    def __init__(self, embed_dim=768, mlp_ratio=4, dropout=0.1):
        super().__init__()
        hidden_dim = int(embed_dim * mlp_ratio)
        self.fc1 = nn.Linear(embed_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, embed_dim)
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, x):
        x = self.fc1(x)
        x = F.gelu(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.dropout(x)
        return x

class TransformerBlock(nn.Module):
    """Transformer编码器块"""
    def __init__(self, embed_dim=768, n_heads=12, mlp_ratio=4, dropout=0.1):
        super().__init__()
        self.norm1 = nn.LayerNorm(embed_dim)
        self.attn = MultiHeadAttention(embed_dim, n_heads, dropout)
        self.norm2 = nn.LayerNorm(embed_dim)
        self.mlp = MLP(embed_dim, mlp_ratio, dropout)
        
    def forward(self, x):
        # 注意力机制 + 残差连接
        x = x + self.attn(self.norm1(x))
        # 前馈网络 + 残差连接
        x = x + self.mlp(self.norm2(x))
        return x

class VisionTransformer(nn.Module):
    """Vision Transformer模型 - 针对MX450优化的轻量级版本"""
    def __init__(self, 
                 img_size=224, 
                 patch_size=16, 
                 in_channels=3, 
                 num_classes=5,
                 embed_dim=384,  # 减小嵌入维度以适应MX450
                 depth=6,        # 减少层数
                 n_heads=6,      # 减少注意力头数
                 mlp_ratio=4, 
                 dropout=0.1):
        super().__init__()
        
        self.num_classes = num_classes
        self.embed_dim = embed_dim
        
        # 图像块嵌入
        self.patch_embed = PatchEmbedding(img_size, patch_size, in_channels, embed_dim)
        n_patches = self.patch_embed.n_patches
        
        # 位置嵌入和类别token
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(torch.zeros(1, n_patches + 1, embed_dim))
        self.dropout = nn.Dropout(dropout)
        
        # Transformer编码器
        self.blocks = nn.ModuleList([
            TransformerBlock(embed_dim, n_heads, mlp_ratio, dropout)
            for _ in range(depth)
        ])
        
        # 分类头
        self.norm = nn.LayerNorm(embed_dim)
        self.head = nn.Linear(embed_dim, num_classes)
        
        # 初始化权重
        self.init_weights()
        
    def init_weights(self):
        """初始化模型权重"""
        # 初始化位置嵌入
        nn.init.trunc_normal_(self.pos_embed, std=0.02)
        nn.init.trunc_normal_(self.cls_token, std=0.02)
        
        # 初始化线性层
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.trunc_normal_(m.weight, std=0.02)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.LayerNorm):
                nn.init.constant_(m.bias, 0)
                nn.init.constant_(m.weight, 1.0)
    
    def forward(self, x):
        batch_size = x.shape[0]
        
        # 图像块嵌入
        x = self.patch_embed(x)  # (batch_size, n_patches, embed_dim)
        
        # 添加类别token
        cls_tokens = self.cls_token.expand(batch_size, -1, -1)
        x = torch.cat([cls_tokens, x], dim=1)  # (batch_size, n_patches + 1, embed_dim)
        
        # 添加位置嵌入
        x = x + self.pos_embed
        x = self.dropout(x)
        
        # 通过Transformer编码器
        for block in self.blocks:
            x = block(x)
        
        # 归一化并提取类别token
        x = self.norm(x)
        cls_token_final = x[:, 0]  # (batch_size, embed_dim)
        
        # 分类
        logits = self.head(cls_token_final)  # (batch_size, num_classes)
        
        return logits
    
    def predict(self, x, return_probs=False):
        """预测函数，返回类别和描述"""
        self.eval()
        with torch.no_grad():
            logits = self.forward(x)
            probs = F.softmax(logits, dim=-1)
            predicted_ids = torch.argmax(logits, dim=-1)
            
            if return_probs:
                return predicted_ids, probs
            else:
                return predicted_ids
    
    def get_model_info(self):
        """获取模型信息"""
        total_params = sum(p.numel() for p in self.parameters())
        trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad)
        
        info = {
            'model_name': 'Vision Transformer (轻量级)',
            'total_parameters': total_params,
            'trainable_parameters': trainable_params,
            'embed_dim': self.embed_dim,
            'num_classes': self.num_classes,
            'class_descriptions': ID_TO_DESCRIPTION
        }
        return info

def create_model(num_classes=5, pretrained=False):
    """创建模型实例"""
    model = VisionTransformer(
        img_size=224,
        patch_size=16,
        num_classes=num_classes,
        embed_dim=384,  # 适合MX450的轻量级配置
        depth=6,
        n_heads=6,
        dropout=0.1
    )
    
    if pretrained:
        print("注意: 当前版本不支持预训练权重加载")
    
    return model

if __name__ == "__main__":
    # 测试模型
    model = create_model()
    
    # 打印模型信息
    info = model.get_model_info()
    print("模型信息:")
    print("-" * 50)
    for key, value in info.items():
        if key != 'class_descriptions':
            print(f"{key}: {value}")
    
    print("\n类别描述:")
    for class_id, desc in info['class_descriptions'].items():
        print(f"类别 {class_id}: {desc}")
    
    # 测试前向传播
    dummy_input = torch.randn(2, 3, 224, 224)
    output = model(dummy_input)
    print(f"\n输入形状: {dummy_input.shape}")
    print(f"输出形状: {output.shape}")
    print(f"输出示例: {output[0].detach().numpy()}")